package com.chatplus.application.listener;

import cn.bugstack.openai.session.Configuration;
import com.chatplus.application.aiprocessor.platform.image.sd.SdImageProcessor;
import com.chatplus.application.common.logging.SouthernQuietLogger;
import com.chatplus.application.common.logging.SouthernQuietLoggerFactory;
import com.chatplus.application.common.util.OkHttpClientUtil;
import com.chatplus.application.common.util.PlusJsonUtils;
import com.chatplus.application.domain.dto.ApiKeyDto;
import com.chatplus.application.domain.dto.SdJobResult;
import com.chatplus.application.domain.entity.draw.SdJobEntity;
import com.chatplus.application.domain.notification.SdJobNotification;
import com.chatplus.application.service.account.UserProductLogService;
import com.chatplus.application.service.draw.SdJobService;
import com.chatplus.application.web.notification.NotificationListener;
import okhttp3.*;
import org.apache.commons.lang3.StringUtils;
import org.springframework.stereotype.Component;

import java.io.IOException;

/**
 * Sd 任务监听
 */
@Component
public class SdJobNotificationListener {

    private final SdImageProcessor sdImageProcessor;
    private final SdJobService sdJobService;
    private final UserProductLogService userProductLogService;

    public SdJobNotificationListener(SdImageProcessor sdImageProcessor, SdJobService sdJobService, UserProductLogService userProductLogService) {
        this.sdImageProcessor = sdImageProcessor;
        this.sdJobService = sdJobService;
        this.userProductLogService = userProductLogService;
    }

    private static final SouthernQuietLogger LOGGER = SouthernQuietLoggerFactory.getLogger(SdJobNotificationListener.class);

    @NotificationListener(notification = SdJobNotification.class)
    public void handleSdJobNotification(SdJobNotification notification) {
        ApiKeyDto apiKeyDto = notification.getApiKeyDto();
        OkHttpClient okHttpClient = OkHttpClientUtil.getOkHttpClient(apiKeyDto.getProxyUrl(), true);
        Request request = new Request.Builder()
                .addHeader("Content-Type", Configuration.APPLICATION_JSON)
                .addHeader("Accept", Configuration.APPLICATION_JSON)
                .url(apiKeyDto.getUrl() + "/sdapi/v1/txt2img")
                .post(RequestBody.create(PlusJsonUtils.toJsonString(notification.getParams()), MediaType.parse(Configuration.APPLICATION_JSON)))
                .build();
        Call call = okHttpClient.newCall(request);
        SdJobEntity sdJobEntity = sdJobService.getById(notification.getSdJobId());
        try (Response response = call.execute()) {
            if (response.isSuccessful() && response.body() != null) {
                SdJobResult result = PlusJsonUtils.parseObject(response.body().string(), SdJobResult.class);
                if (result == null || result.getImages() == null || result.getImages().length == 0) {
                    LOGGER.message("SD图片生成失败").context("response", response).error();
                    sdJobEntity.setProgress(-1);
                    return;
                }
                String url = sdImageProcessor.saveToOss(null, result.getImages()[0]);
                if (StringUtils.isNotEmpty(url)) {
                    sdJobEntity.setImgUrl(url);
                    sdJobEntity.setProgress(100);
                    sdImageProcessor.notifyUpdateTask(sdJobEntity.getId());
                }
            } else {
                LOGGER.message("SD图片请求失败").context("response", response).error();
                userProductLogService.adminAddProductToUser(sdJobEntity.getUserId(), 0, 1);
                sdJobEntity.setProgress(-1);
            }
        } catch (IOException e) {
            sdJobEntity.setProgress(-1);
            userProductLogService.adminAddProductToUser(sdJobEntity.getUserId(), 0, 1);
            LOGGER.message("SD图片请求失败").exception(e).error();
        } finally {
            sdJobService.updateById(sdJobEntity);
        }
    }
}
